import numpy as np


def cone_plane_cross(Q, plane):
    """
    求解圆锥与平面的交点
    cone: 圆锥
    plane: 平面
    """
    # 提取平面方程的参数
    a, b, c, d = plane
    if len(Q.shape) == 2:
        Q = Q[None, :, :]

    # 计算交线的齐次矩阵 Q'
    A_prime = Q[..., 0, 0] * c**2 - 2 * a * c * Q[..., 0, 2] + a**2 * Q[..., 2, 2]
    E_prime = Q[..., 1, 1] * c**2 - 2 * b * c * Q[..., 1, 2] + b**2 * Q[..., 2, 2]
    B_prime = (
        Q[..., 0, 1] * c**2
        - a * c * Q[..., 1, 2]
        - b * c * Q[..., 0, 2]
        + a * b * Q[..., 2, 2]
    )
    D_prime = (
        Q[..., 0, 3] * c**2
        - a * c * Q[..., 2, 3]
        - d * c * Q[..., 0, 2]
        + a * d * Q[..., 2, 2]
    )
    G_prime = (
        Q[..., 1, 3] * c**2
        - b * c * Q[..., 2, 3]
        - d * c * Q[..., 1, 2]
        + b * d * Q[..., 2, 2]
    )
    J_prime = Q[..., 3, 3] * c**2 - 2 * d * c * Q[..., 2, 3] + d**2 * Q[..., 2, 2]

    return np.array(
        [
            [A_prime, B_prime, D_prime],
            [B_prime, E_prime, G_prime],
            [D_prime, G_prime, J_prime],
        ]
    ).T.squeeze()


def ransac_ellipse_fit(x, y, t=6, k=100, th=10):
    idx = np.arange(len(x))
    sample = np.random.choice(idx, t, replace=False)
    best_num = 0
    best_choice = None
    x_homo = np.pad((x, y), ((0, 1), (0, 0)), constant_values=1)
    for _ in range(k):
        C = ellipse_fit(x[sample], y[sample])
        err = np.abs(x_homo * (C @ x_homo)).sum(0) / x_homo[2] ** 2
        inliers = err < th
        if inliers.sum() > best_num:
            best_num = inliers.sum()
            best_choice = inliers
            if best_num > t:
                sample = idx[inliers]
            else:
                sample = np.random.choice(idx, t, replace=False)
        else:
            sample = np.random.choice(idx, t, replace=False)
    if best_choice is None:
        return None
    else:
        return ellipse_fit(x[best_choice], y[best_choice])


def ellipse_fit(x, y):
    assert x.shape[-1] > 6
    A = np.stack((x * x, 2 * x * y, y * y, 2 * x, 2 * y, np.ones_like(x)), axis=-1)
    _, _, V = np.linalg.svd(A)
    X = V[:, -1]
    # 计算不确定度
    v = np.einsum("nmi,ni->nm", A, X)
    D = (v * v).sum() / (v.size - 6) * (V * V).sum(axis=1)
    return homogenous_ellipse(*X.T), np.sqrt(D)


def homogenous_ellipse(A, B, C, D, E, F):
    return np.array([[A, B, D], [B, C, E], [D, E, F]]).T


def from_ellipse_args(x, y, a, b, theta):
    idx = np.linspace(0, 2 * np.pi, 200)
    x = (
        x[..., None]
        + a[..., None] * np.cos(idx) * np.cos(theta[..., None])
        - b[..., None] * np.sin(idx) * np.sin(theta[..., None])
    )
    y = (
        y[..., None]
        + a[..., None] * np.cos(idx) * np.sin(theta[..., None])
        + b[..., None] * np.sin(idx) * np.cos(theta[..., None])
    )
    return np.stack((x, y), axis=1)


def center_ellipse(C):
    if len(C.shape) == 2:
        C = C[None, :, :]
    x_cnt = (C[..., 0, 1] * C[..., 1, 2] - C[..., 1, 1] * C[..., 0, 2]) / (
        C[..., 0, 0] * C[..., 1, 1] - C[..., 0, 1] * C[..., 0, 1]
    )
    y_cnt = (C[..., 0, 1] * C[..., 0, 2] - C[..., 0, 0] * C[..., 1, 2]) / (
        C[..., 0, 0] * C[..., 1, 1] - C[..., 0, 1] * C[..., 0, 1]
    )
    return x_cnt, y_cnt


def radius_ellipse(C):
    if len(C.shape) == 2:
        C = C[None, :, :]
    a, b, c, d, e, f = (
        C[..., 0, 0],
        C[..., 0, 1],
        C[..., 1, 1],
        C[..., 0, 2],
        C[..., 1, 2],
        C[..., 2, 2],
    )
    val1 = (
        2
        * (a * e * e - 2 * b * d * e + c * d * d - (a * c - b * b) * f)
        / (a * c - b * b)
        / (a + c - np.sqrt((a - c) ** 2 + 4 * b * b))
    )
    val2 = (
        2
        * (a * e * e - 2 * b * d * e + c * d * d - (a * c - b * b) * f)
        / (a * c - b * b)
        / (a + c + np.sqrt((a - c) ** 2 + 4 * b * b))
    )
    ind = (val1 >= 0) & (val2 >= 0)
    r1 = np.zeros_like(val1) * np.nan
    r2 = np.zeros_like(val2) * np.nan
    r1[ind] = np.sqrt(val1[ind])
    r2[ind] = np.sqrt(val2[ind])
    return r1, r2


def ellipse_args(C):
    if len(C.shape) == 2:
        C = C[None, :, :]
    x = center_ellipse(C)
    a, b, c, d, e, f = (
        C[..., 0, 0],
        C[..., 0, 1],
        C[..., 1, 1],
        C[..., 0, 2],
        C[..., 1, 2],
        C[..., 2, 2],
    )
    val1 = (
        2
        * (a * e * e - 2 * b * d * e + c * d * d - (a * c - b * b) * f)
        / (a * c - b * b)
        / (a + c - np.sqrt((a - c) ** 2 + 4 * b * b))
    )
    val2 = (
        2
        * (a * e * e - 2 * b * d * e + c * d * d - (a * c - b * b) * f)
        / (a * c - b * b)
        / (a + c + np.sqrt((a - c) ** 2 + 4 * b * b))
    )
    ind = (val1 >= 0) & (val2 >= 0)
    r1 = np.zeros_like(val1) * np.nan
    r2 = np.zeros_like(val2) * np.nan
    theta = np.zeros_like(val1) * np.nan
    r1[ind] = np.sqrt(val1[ind])
    r2[ind] = np.sqrt(val2[ind])
    theta[ind] = np.arctan2(
        2 * b[ind],
        a[ind] - c[ind] - np.sqrt((a[ind] - c[ind]) ** 2 + 4 * b[ind] * b[ind]),
    )
    return *x, r1, r2, theta


def ellipse_args_inv(x, y, a, b, theta):
    """由椭圆参数计算椭圆一般方程"""
    A = a * a * np.sin(theta) ** 2 + b * b * np.cos(theta) ** 2
    B = (b * b - a * a) * np.sin(theta) * np.cos(theta)
    C = a * a * np.cos(theta) ** 2 + b * b * np.sin(theta) ** 2
    D = -A * x - B * y / 2
    E = -B * x / 2 - C * y
    F = A * x * x + B * x * y + C * y * y - a * a * b * b
    return homogenous_ellipse(A, B, C, D, E, F)


def center_uncertainty(C, U):
    """
    C : n x 3 x 3，齐次矩阵
    U : n x 6 : 齐次矩阵的六个独立参数
    """
    if len(C.shape) == 2:
        C = C[None, :, :]
    if len(U.shape) == 1:
        U = U[None, :]
    a, b, c, d, e, f = (
        C[..., 0, 0],
        C[..., 0, 1],
        C[..., 1, 1],
        C[..., 0, 2],
        C[..., 1, 2],
        C[..., 2, 2],
    )
    D0 = b * b - a * c
    u_x_0 = (
        1
        / D0
        * np.sqrt(
            ((c * d - b * e) / D0 * U[..., 0]) ** 2
            + ((e * b * b - a * e * c - 2 * b * c * d) / D0 * U[..., 1]) ** 2
            + ((d * b * b - a * b * e) / D0 * U[..., 2]) ** 2
            + (c * U[..., 3]) ** 2
            + (b * U[..., 4]) ** 2
        )
    )
    u_y_0 = (
        1
        / D0
        * np.sqrt(
            ((b * b * e - c * b * d) / D0 * U[..., 0]) ** 2
            + ((d * b * b - a * c * d - 2 * a * b * e) / D0 * U[..., 1]) ** 2
            + ((a * e - b * d) / D0 * U[..., 2]) ** 2
            + (b * U[..., 3]) ** 2
            + (a * U[..., 4]) ** 2
        )
    )
    return u_x_0, u_y_0


def axis_uncertainty(C, U, r1, r2):
    """
    C : n x 3 x 3，齐次矩阵
    U : n x 6 : 齐次矩阵的六个独立参数
    """
    if len(C.shape) == 2:
        C = C[None, :, :]
    if len(U.shape) == 1:
        U = U[None, :]
    a, b, c, d, e, f = (
        C[..., 0, 0, None],
        C[..., 0, 1, None],
        C[..., 1, 1, None],
        C[..., 0, 2, None],
        C[..., 1, 2, None],
        C[..., 2, 2, None],
    )
    D0 = b * b - a * c
    Na_b = a * e * e + c * d * d + b * f * f - 2 * b * d * e - a * c * f
    Dabc = np.sqrt((a - c) ** 2 + 4 * b * b)
    D1 = Dabc - a - c
    D2 = -Dabc - a - c
    p_N_a_b = np.concatenate(
        (
            e * e * -c * f,
            f * f * -2 * d * e,
            d * d - a * f,
            2 * c * d - 2 * b * e,
            2 * a * e - 2 * b * d,
            2 * b * f - a * c,
        ),
        axis=2,
    )
    p_D_a = np.pad(
        (
            -c * D1 + D0 * ((a - c) / Dabc - 1),
            2 * b * Dabc + D0 * 4 * b / Dabc,
            -a * D1 + D0 * ((c - a) / Dabc - 1),
        ),
        ((0, 0), (0, 0), (0, 3)),
        constant_values=0,
    )
    p_D_b = np.pad(
        (
            -c * D2 + D0 * (-(a - c) / Dabc - 1),
            -2 * b * Dabc - D0 * 4 * b / Dabc,
            -a * D2 + D0 * (-(c - a) / Dabc - 1),
        ),
        ((0, 0), (0, 0), (0, 3)),
        constant_values=0,
    )
    u_a = (
        1
        / r1
        * np.sqrt(
            np.power((p_N_a_b * D0 * D1 - p_D_a * Na_b) / (D0 * D1) ** 2 * U, 2).sum(
                axis=2
            )
        )
    )
    u_b = (
        1
        / r2
        * np.sqrt(
            np.power((p_N_a_b * D0 * D2 - p_D_b * Na_b) / (D0 * D2) ** 2 * U, 2).sum(
                axis=2
            )
        )
    )
    return u_a, u_b
